/* 
 * Based on Lowe, "Distinctive Image Features form Scale-Invariant Keypoints",
 * International Journal of Computer Vision, 2004
 */

#include<math.h>
#include "sc.h"

#define SLICES  8
#define OCTAVES 1
#define MEASURE
typedef float descriptor_t[128];

typedef struct
{
  int x, y;
} point_t;

PixelToPixelOp()
itof(in stream unsigned char *i,
     out stream float *o)
{
  *o = *i;
}

NeighbourhoodToPixelOp()
gaussy(in stream float i[-9..9][0],
       in double *sigma,
       out stream float *o)
{
  #define SIZE 9
  int k;
  float kernel[SIZE];
  float acc=0;

  init {
    extern double sqrt(double x);
    extern double exp(double x);
    float M_PI=3.14159265358979323846264338327950288419;
    float sum=0;
    
    for (k=-SIZE/2; k <= SIZE/2; k++)
      sum += kernel[k+SIZE/2] = 1/(sqrt(2*M_PI)**sigma)*exp(-(k*k)/(2**sigma**sigma));

    for (k=-SIZE/2; k <= SIZE/2; k++)
      kernel[k+SIZE/2] /= sum;
  };

  for (k=-SIZE/2; k <= SIZE/2; k++)
    acc += i[k][0] * kernel[k+SIZE/2];

  *o = acc;
}

NeighbourhoodToPixelOp()
gaussx(in stream float i[0][-9..9],
       in double *sigma,
       out stream float *o)
{
  #define SIZE 9
  int k;
  float kernel[SIZE];
  float acc=0;

  init {
    extern double sqrt(double x);
    extern double exp(double x);
    float M_PI=3.14159265358979323846264338327950288419;
    float sum=0;
    
    for (k=-SIZE/2; k <= SIZE/2; k++)
      sum += kernel[k+SIZE/2] = 1/(sqrt(2*M_PI)**sigma)*exp(-(k*k)/(2**sigma**sigma));

    for (k=-SIZE/2; k <= SIZE/2; k++)
      kernel[k+SIZE/2] /= sum;
  };

  for (k=-SIZE/2; k <= SIZE/2; k++)
    acc += i[0][k] * kernel[k+SIZE/2];

  *o = acc;
}

PixelToPixelOp()
subtract(in stream float *i1,
         in stream float *i2,
         out stream float *o)
{
  *o = *i1-*i2;
}

NeighbourhoodToPixelOp()
findextrema(in stream float i1[-1..1][-1..1],
            in stream float i2[-1..1][-1..1],
            in stream float i3[-1..1][-1..1],
            out stream int *e,
            in int *dummy)
{
  extern float fabsf(float x);
  float v = i2[0][0];

  *e = 0;

  /* p11, "... all extrema with a value of |D(x)| less than 0.03 were discarded " */
  if (fabsf(v) > 0.008*256)
  {  
    int x, y;
    int min=0, max=0;

    for (y=-1; y < 2; y++)
      for (x=-1; x < 2; x++)
      {
        min = min + (v < i1[y][x]);
        if (x || y) min = min + (v < i2[y][x]);
        min = min + (v < i3[y][x]);

        max = max + (v > i1[y][x]);
        if (x || y) max = max + (v > i2[y][x]);
        max = max + (v > i3[y][x]);
      }

    if (min == 26 || max == 26)
      *e = 1;
  }
}

AnisoPixelToPixelOp()
genpoints(in stream float *i,
          out stream point_t *p)
{
  (*p).x = x;
  (*p).y = y;
}

PixelFilterOp()
filterextrema(in stream int *e,
              in stream point_t *p_in,
              out stream point_t *p_out)
{
  if (*e && (*p_in).x > 8 && (*p_in).y > 8 && (*p_in).x < 1024-8 && (*p_in).y < 1024-8)
  {
/*    printf("filtextr: (%d, %d)\n", (*p_in).x, (*p_in).y);*/

    *p_out = *p_in;
    push();
  }
}

PixelLookupOp()
hessianproportion(in stream point_t *p,
                  lookup float i[][],
                  out stream float *h)
{
  /*
   * p12, "... to check that the ratio of principal curvatures is
   * below some threshold r, we only need to check
   * tr(H)^2/Det(H) < (r+1)^2/r "
   */

  float dxx = i[(*p).y][(*p).x-1] - 2*i[(*p).y][(*p).x] + i[(*p).y][(*p).x+1];
  float dyy = i[(*p).y-1][(*p).x] - 2*i[(*p).y][(*p).x] + i[(*p).y+1][(*p).x];
  float dxy = i[(*p).y+1][(*p).x+1] + i[(*p).y][(*p).x] - i[(*p).y+1][(*p).x] - i[(*p).y][(*p).x+1];
  
  float tr = dxx+dyy;
  float det = dxx*dyy-dxy*dxy;
  *h = tr*tr/det;  

/*  printf("hessprop: (%d, %d) -> %f\n", (*p).x, (*p).y, *h); */
}

PixelFilterOp()
filterproportion(in stream float *h,
                 in stream point_t *p_in,
                 out stream point_t *p_out)
{
  /* P12, "The experiments in this paper use a value of r = 10"
   * -> (r+1)^2/r = 12.1
   */
  
  if (*h < 12.1)
  {
/*    printf("filtprop: (%d, %d) -> %f\n", (*p_in).x, (*p_in).y, *h); */
    
    *p_out = *p_in;
    push();
  }
}

NeighbourhoodToPixelOp()
orientation(in stream float i[-1..1][-1..1],
            out stream float *m,
            out stream float *t,
            in int *dummy)
{
  extern double atan2(double y, double x);
  extern double sqrt(double x);
  
  float dx = i[0][1]-i[0][-1];
  float dy = i[1][0]-i[-1][0];
  
  *m = sqrt(dx*dx + dy*dy);
  *t = atan2(dy,dx);
}

PixelLookupOp()
findorientation(in stream point_t *p,
                in double *s,
                lookup float m[][],
                lookup float t[][],
                out stream float *o)
{
  float hist[36];
  float kernel[7][7];
  int kx, ky, h, hmax=0;
  
  init {
    /* p13, "... Gaussian-weighted circular window with a \sigma that is
     * 1.5 times that of the scale of the keypoint" */

    extern double pow(double q, double r);
    extern double exp(double q);
    float M_PI=3.14159265358979323846264338327950288419;
    float sigma = *s * 1.5;
    float sum = 0;

    for (ky=-3; ky < 4; ky++)
      for (kx=-3; kx < 4; kx++)
        sum += kernel[ky+3][kx+3] = 1/(2*M_PI*sigma*sigma)*exp(-(kx*kx+ky*ky)/(2*sigma*sigma));
  
    for (ky=-3; ky < 4; ky++)
      for (kx=-3; kx < 4; kx++)
        kernel[ky+3][kx+3] /= sum;
  };

  memset(hist, 0, 36*sizeof(float));

  for (ky=-3; ky < 4; ky++)
    for (kx=-3; kx < 4; kx++)
    {
       /* p13, "The orientation histogram has 36 bins" */
       int bin = (int)((t[(*p).y+ky][(*p).x+kx]+M_PI)/(2*M_PI)*(36-0.0001));
       
       /* p13, "Each sample added to the histogram is weighted by its gradient
        * magnitude and by a Gaussian ..." */    
       hist[bin] = hist[bin] + m[(*p).y+ky][(*p).x+kx] * kernel[ky+3][kx+3];
    }
  
  for (h=1; h < 36; h++)
    if (hist[h] > hist[hmax])
      hmax = h;

  *o = hmax/36.0*2*M_PI-M_PI;

/*  printf("findorient: (%d, %d) -> %f (%f)\n", (*p).x, (*p).y, *o, hist[hmax]); */
}

PixelLookupOp()
generatedescriptors(in stream point_t *p,
                    in stream float *o,
                    lookup float m[][],
                    lookup float t[][],
                    out stream descriptor_t *desc)
{
  int dx, dy, kx, ky;
  float kernel[16][16];

  init {
    extern double pow(double q, double r);
    extern double exp(double q);
    float M_PI=3.14159265358979323846264338327950288419;
    float sigma = 8;
    float sum = 0;

    /* p15, "A Gaussian weighting function with \sigma equal to one half the width of
     * the discriptor window is used to assign a weight to the magnitude of each point */
    for (ky=-8; ky < 8; ky++)
      for (kx=-8; kx < 8; kx++)
        sum += kernel[ky+8][kx+8] = 1/(2*M_PI*sigma*sigma)*exp(-((kx+0.5)*(kx+0.5)+(ky+0.5)*(ky+0.5))/(2*sigma*sigma));
      
    for (ky=-8; ky < 8; ky++)
      for (kx=-8; kx < 8; kx++)
        kernel[ky+8][kx+8] /= sum;
  };  

/*  printf("generatedesc: (%d, %d)\n", (*p).x, (*p).y); */

  memset(*desc, 0, sizeof(descriptor_t));
  
  /* p15, "... samples are the accumulated into orientation histograms summarizing
   * the contents over 4x4 subregions..." */
  for (dy=-2; dy < 2; dy++)
    for (dx=-2; dx < 2; dx++)
      for (ky=0; ky < 4; ky++)
        for (kx=0; kx < 4; kx++)
        {
          /* p15, "... eight directions for each orientation histogram... " */
          int bin = (int)((t[(*p).y+dy*4+ky][(*p).x+dx*4+kx]-*o)/(2*M_PI)*8);
          int el = (ky*4+kx)*8+((bin+20)%8);

          (*desc)[el] += m[(*p).y+dy*4+ky][(*p).x+dx*4+kx]*kernel[dy*4+ky+8][dx*4+kx+8];
        }
}

PixelToPixelOp()
normalize(in stream descriptor_t *i,
          out stream descriptor_t *o)
{
  float total=0;
  int d;
  
  for (d=0; d < 128; d++)
    total = total + (*i)[d];

  for (d=0; d < 128; d++)
  {
    /* p16, "... the vector is normalized to unit length." */
    (*o)[d] = (*i)[d] / total;

    /* p16, "... thresholding the values in the unit feature vector to each be no
     * larger than 0.2 ... " */
    if ((*o)[d] > 0.2)
      (*o)[d] = 0.2;
  }
  
  total = 0;
  for (d=0; d < 128; d++)
    total += (*i)[d];

  for (d=0; d < 128; d++)
  {
    /* p16, "... and then renormalizing to unit length." */
    (*o)[d] = (*i)[d] / total;
  }
}

PixelToPixelOp()
print(in stream descriptor_t *i)
{
  int d;

  printf("{ ");
  for (d=0; d < 128; d++)
    printf("%f ", (*i)[d]);
  printf("}\n");
}

PixelToPixelOp()
ftoi(in stream float *i,
     in int *dummy,
     out stream unsigned char *o)
{
  *o = *i;
}

PixelToPixelOp()
ftoi2(in stream float *i,
      in int *dummy,
      out stream unsigned char *o)
{
  extern float fabsf(float x);

  *o = 64 + (fabsf(*i)>0.008*256)*128 + *i*4;
}

PixelToGlobalOp(in int *width=scint(1024), in int *height=scint(1024))
drawpoints(in stream point_t *p,
           out stream unsigned char **o)
{
  int i;

/*  printf("drawpoints: (%d, %d)\n", (*p).x, (*p).y); */

  for (i=-8; i < 9; i++)
  {
    o[(*p).y+i][(*p).x-8] = 255;
    o[(*p).y+i][(*p).x+8] = 255;
    o[(*p).y-8][(*p).x+i] = 255;
    o[(*p).y+8][(*p).x+i] = 255;
  }
}

PixelToPixelOp()
overlay(in stream unsigned char *original,
        in stream unsigned char *overlay,
        out stream unsigned char *result)
{
  if (*overlay)
    *result = *overlay;
   else
    *result = *original;
}

MAINTYPE MAIN(int argc, char **argv)
{
  int o, s, sl;
  double k = pow(2, 1.0/SLICES), sigma;
  STREAM image, gss[OCTAVES][SLICES+3], dog[OCTAVES][SLICES+2], keypoints, ovl, disp;
  scInit(argc, argv);
  _allmulti(1);
#ifndef MEASURE
  capture(&image);
#else
  generate(scint(1024), &image);
#endif
  itof(&image, &gss[0][0]);

  for (o=0; o < OCTAVES; o++)
  {
    /* Build gaussian scalespace for this octave */
    for (s=1, sigma=k; s < SLICES+3; s++, sigma*=k)
    {
      /* gauss(&gss[o][0], &sigma, &gss[o][s]); */
      gaussx(&gss[o][0], &sigma, &gss[o][s]);
      gaussy(&gss[o][s], &sigma, &gss[o][s]);
    }      

    /* Build difference of gaussian stack */
    for (s=0; s < SLICES+2; s++)
      subtract(&gss[o][s+1], &gss[o][s], &dog[o][s]);

    /* Subsample to next octave in piramid */
    /*
    if (o < OCTAVES-1)
      subsample(&gss[o][SLICES+1], &gss[o+1][0], 2);
    */
    
    for (sl=1; sl < SLICES+1; sl++)
    {
      STREAM p, m, t, extrema, hthr, orientation;
    
      s = sl;

      /* Find local extrema in 3D */
      findextrema(&dog[o][s-1], &dog[o][s], &dog[o][s+1], &extrema, &sl);

      /* Keep only extrema */
      genpoints(&extrema, &p);
      filterextrema(&extrema, &p, &p);

      /* Calculate proportion of hessian */
      hessianproportion(&p, &gss[o][s], &hthr);
      
      /* Keep only keypoints which satisfy a threshold on the hessian proportion */
      filterproportion(&hthr, &p, &p);

      /* Get gradient magnitude and angle */
      orientation(&gss[o][s], &m, &t, &sl);

      /* Find principal orientation */
      findorientation(&p, scdbl(pow(2, (double)s/SLICES)), &m, &t, &orientation);

      /* Generate keypoint descriptors */
      generatedescriptors(&p, &orientation, &m, &t, &keypoints);
      
      /* Normalize to unit length */
      normalize(&keypoints, &keypoints);
#ifndef MEASURE
      /* Some output */
      print(&keypoints);

      write(overlay(ftoi2(&dog[o][s], &sl, NULL),
                    drawpoints(&p, NULL),
                    NULL),
            &s);
#else
      discard(&keypoints);
#endif
      _finalize(&keypoints);
      _finalize(&orientation);
      _finalize(&m);
      _finalize(&t);
      _finalize(&hthr);
      _finalize(&p);
      _finalize(&extrema);
    }

    for (s=0; s < SLICES+2; s++)
      _finalize(&dog[o][s]);

    /* Don't need the gaussian scalespace anymore */
    for (s=0; s < SLICES+3; s++)
      _finalize(&gss[o][s]);
  }

  _finalize(&image);

  return scExit();
}
